# -*- coding: utf-8 -*-
"""
Created on Mon Sep 09 11:19:14 2013

@author: jkwong
"""

# PBAR_ExamineZspecGainDatasets.py

import PBAR_Zspec
import numpy as np
import copy
import csv
from scipy.optimize import curve_fit
from scipy.stats import tmean
# PMT curve fit
def PMT_Curve(x, a, b):
    return(a* np.power(x, b))

# List of good and bad detectors
badZspecList = np.array([1,2,3,4,5,6,7,8,20,26,31,33,34,38,39,40,44,53,56,62,68,76,80,125,126,127,128,129,130,131,132,133,134,135,136])
goodZspecList = np.array([i for i in np.arange(1, 137) if (i not in badZspecList)])

badZspecList -= 1
goodZspecList -= 1

baseDir = r'C:\Users\jkwong\Documents\Work\PBAR\data'


# create list of dataset names
# not including ea91 because bad
datasetNameList = np.array(['ea%d'%i for i in xrange(81, 96) if i != 91])
filenameList = np.array([(i+'.csv') for i in datasetNameList])
fullFilenameList = np.array([os.path.join(baseDir,i+'.csv') for i in datasetNameList])
dat = PBAR_Zspec.ReadZspec(fullFilenameList)

# Important dataset information
voltageOffset = np.array([0, -9, -6, 0, 9, 21, 36, -3, 0, 3, 6, 12, 0, 5])
epochTimeList = np.array([1378493600, 1378494983, 1378495716, 1378496415, 1378497181, 1378498125, 1378498872, 1378500092, 1378500876,\
    1378501619, 1378503186, 1378503992, 1378505158, 1378507618])
epochTimeList0 = epochTimeList - epochTimeList[0]
acqTime = 300.0

# read in voltage
voltageDat = np.genfromtxt(r'C:\Users\jkwong\Documents\Work\PBAR\data\voltage_20130906_121626.txt', delimiter = ',')

voltageSet = voltageDat[:,3]
voltageMeasured = voltageDat[:,4]


binBaseline = 100.
rateRange = [0.5, 2.0] 
rateLook = 1.0
binBounds = [20, 200]

gainShiftAll = np.zeros((len(dat), dat[0].shape[1]))
gainShift0All = np.zeros((len(dat), dat[0].shape[1]))


for i in xrange(len(dat)):
    #    print dat[0].shape
    binArray = np.arange(dat[i].shape[0])
    
    # fit an exponential
    #    pfitAll = np.zeros((dat[0].shape[1], 2))
    binMeanAll = np.zeros(dat[i].shape[1])
    
    for ch in xrange(dat[i].shape[1]):
        cutt = (binArray > binBounds[0]) & (binArray < binBounds[1]) & \
        (dat[i][:,ch]/acqTime > rateRange[0]) & (dat[i][:,ch]/acqTime < rateRange[1])
        if any(cutt):
            temp = np.polyfit(binArray[cutt], np.log(dat[i][cutt,ch]/acqTime),1)
            pfit0 = [temp[0], np.exp(temp[1])]
            binMean = (1.0/pfit0[0]) * np.log(rateLook/pfit0[1])
#            print 'a'
        else:
            binMean = 0.0
        binMeanAll[ch] = binMean
    #        print i+1, binMean
    gainShift0All[i,:] = binMeanAll / binBaseline

for i in xrange(gainShift0All.shape[0]):
    gainShiftAll[i,:]= gainShift0All[i,:] / gainShift0All[0,:]


# First we have to remove the gain shift due to temperature buy using the deltaV = 0 datasets

gainShiftCorrectedAll =  copy.copy(gainShiftAll)
# apply correction gain shift due to temperature
cut = voltageOffset == 0    
x = epochTimeList0[cut]
y = gainShiftAll[cut,:]

# Correct each channel separately
pfitGainShiftVsTimeList = np.zeros((136, 2))
for i in xrange(136):
    pfit = np.polyfit(x, log(y[:,i]), 1)
    pfitGainShiftVsTimeList[i,:] = pfit
    # correct gain
    gainShiftCorrectedAll[:,i] = gainShiftAll[:,i] * np.exp(np.polyval(pfit, 0)) / np.exp(np.polyval(pfit, epochTimeList0))


# Fit PMT voltage gain curves, one PMT at a time
x = voltageOffset
y = gainShiftCorrectedAll

pfitGainShiftVsVoltageList = np.zeros((136, 2))
poptGainShiftVsVoltageList = np.zeros((136, 2))
for i in xrange(136):
    cut = y[:,i] > 0.0
    try:
        pfit = np.polyfit(x[cut], log(y[cut,i]), 1)
    except:
        pfit = np.array([-1, -1])
    pfitGainShiftVsVoltageList[i,:] = pfit
    try:
        startingParam = [2., 20.]
        popt, pcov = curve_fit(PMT_Curve, (x[cut] + voltageSet[i])/1000.0, y[cut,i], p0 = startingParam)
#        poptGainShiftVsVoltageList[i,0] = popt[0]
        poptGainShiftVsVoltageList[i,0] = popt[0] / 1000.0 ** popt[1]     
        poptGainShiftVsVoltageList[i,1] = popt[1]
    except:
        poptGainShiftVsVoltageList[i,:] = np.array([-1,-1])

gainCorrectionCoefficient = 1/pfitGainShiftVsVoltageList[:,0]
# correct the channels that have nan values
cutt = ~isnan(gainCorrectionCoefficient) & ~isinf(gainCorrectionCoefficient) & (gainCorrectionCoefficient > 0)
gainCorrectionCoefficientAverage = gainCorrectionCoefficient[cutt].mean()
gainCorrectionCoefficient[~cutt] = gainCorrectionCoefficientAverage

fid = open('GainCalibrationCoefficients.txt', 'wb')
csvWriterObj = csv.writer(fid)
csvWriterObj.writerow(gainCorrectionCoefficient)
fid.close()

# correct the other aray
gainCorrectionCoefficient2 = poptGainShiftVsVoltageList[:,1]
cutt = ~isnan(gainCorrectionCoefficient2) & ~isinf(gainCorrectionCoefficient2) & (gainCorrectionCoefficient2 > 0)
gainCorrectionCoefficient2Average = gainCorrectionCoefficient2[cutt].mean()
gainCorrectionCoefficient2[~cutt] = gainCorrectionCoefficient2Average


fid = open('GainCalibrationCoefficients2.txt', 'wb')
csvWriterObj = csv.writer(fid)
csvWriterObj.writerow(gainCorrectionCoefficient2)
fid.close()


# PLOTS

# VERSUS TIME

# plot gainshift versus time for delta v = 0 
cut = voltageOffset == 0
plt.figure()
plt.grid()
detectorIndex = 57
plot(epochTimeList[cut] - epochTimeList[cut][0], gainShiftAll[cut,detectorIndex], 'xk', markersize = 12, mew = 2)

plt.xlabel('Time Since First Dataset (sec)')
plt.ylabel('Gain Shift')
plt.title('Detector #%d' %(detectorIndex+1))



# plot gainshift versus time for delta v = 0 
plt.figure()
detectorIndex = 60
plt.plot(epochTimeList, gainShiftAll[:,detectorIndex], '.k')

plt.xlabel('Voltage Offset')
plt.ylabel('Gain Shift')
plt.title('Detector #%d' %i)

# GAIN SHIFT VERSUS TIME FOR V = 0
plt.figure()
plt.grid()
cut = voltageOffset == 0
y = gainShift0All[cut,:]
epochTimeListShortList = epochTimeList0[cut]

for i in xrange(y.shape[0]):
    cut2 = ~isnan(y[i,:])
    plt.plot(np.arange(1,137)[cut2], y[i,cut2]/y[0,cut2], marker = 'x', label = '%d sec' %epochTimeListShortList[i])

plt.xlabel('Detector Number')
plt.ylabel('Gain Shift')
plt.legend()
plt.ylim((0.7, 1.3))


# Gain correction parameters as function of detector number
plt.figure()
plt.grid()
cut = ~isnan(gainCorrectionCoefficient)
plot(np.arange(1, 137)[cut], gainCorrectionCoefficient[cut], 'xk')
plt.xlabel('Detector Number')
plt.ylabel('Gain Correction Parameter')


# Gainshift as a function of voltage; good detectors
plt.figure()
for i in xrange(136):
    if i in goodZspecList:
        plt.plot(voltageOffset, gainShiftAll[:,i])


# GAIN SHIFT VERSUS VOLTAGE WITH FITS

plt.figure()
plt.grid()

i = 68

plt.plot(voltageOffset, gainShiftAll[:,i], '.k', label = 'Uncorrected', markersize = 15, alpha = 0.5)
plt.plot(voltageOffset, gainShiftCorrectedAll[:,i], 'xr', label = 'Corrected', markersize = 15, alpha = 0.5)

xarray = np.arange(-10,40)
yarray = np.exp(np.polyval(pfitGainShiftVsVoltageList[i,:], xarray))
plt.plot(xarray, yarray, label = \
    'Best fit to corrected, a*exp(b*v), a = %3.3e, b = %3.3f' %(exp(pfitGainShiftVsVoltageList[i,1]), pfitGainShiftVsVoltageList[i,0]))

plt.plot(xarray, PMT_Curve(xarray + voltageSet[i], poptGainShiftVsVoltageList[i,0], poptGainShiftVsVoltageList[i,1]), \
    label = 'Best fit to corrected, a*v^b, a = %3.3e, b = %3.3f' %(poptGainShiftVsVoltageList[i,0], poptGainShiftVsVoltageList[i,1]))

plt.xlabel('Voltage Offset')
plt.ylabel('Gain Shift')
plt.title('Detector #%d, V = %d' %(i+1, voltageSet[i]))
plt.legend()
plt.yscale('log')



# plot gainshift as a function of voltage, single detector

for i in goodZspecList:
    
    try:
        
        plt.figure()
        plt.grid()
    
        #i = 60
        plt.plot(voltageOffset, gainShiftAll[:,i], '.k', label = 'Uncorrected', markersize = 15, alpha = 0.5)
        plt.plot(voltageOffset, gainShiftCorrectedAll[:,i], 'xr', label = 'Corrected', markersize = 15, alpha = 0.5)
        
        xarray = np.arange(-10,40)
        yarray = np.exp(np.polyval(pfitGainShiftVsVoltageList[i,:], xarray))
        plt.plot(xarray, yarray, label = \
            'Best fit to corrected, a*exp(b*v), a = %3.3e, b = %3.3f' %(exp(pfitGainShiftVsVoltageList[i,1]), pfitGainShiftVsVoltageList[i,0]))
        
        plt.plot(xarray, PMT_Curve(xarray + voltageSet[i], poptGainShiftVsVoltageList[i,0], poptGainShiftVsVoltageList[i,1]), \
            label = 'Best fit to corrected, a*v^b, a = %3.3e, b = %3.3f' %(poptGainShiftVsVoltageList[i,0], poptGainShiftVsVoltageList[i,1]))
        
        plt.xlabel('Voltage Offset')
        plt.ylabel('Gain Shift')
        plt.title('Detector #%d, V = %d' %(i, voltageSet[i]))
        plt.legend()
        plt.yscale('log')
        savefig('GainCalibration_%d' %(i+1))
        plt.close()
    except:
        print("Skipping")


# EXAMINE FIT PARAMETERS

# Examine the fit parameters
# fit 1 parameters

plt.figure()
plt.grid()
cut = goodZspecList
plot(np.arange(1,137), poptGainShiftVsVoltageList[:,1], '.k')
plot(np.arange(1,137)[cut], poptGainShiftVsVoltageList[cut,1], '.r')

# calculate the trimmed mean
cutt = poptGainShiftVsVoltageList[cut,1] > 0
plot([0, 136], poptGainShiftVsVoltageList[cut,1][cutt].mean() * np.ones(2), '-r')

plt.xlabel('Detector Number')
plt.ylabel("PMT Coefficient")

# Examine the fit parameters
# fit 2 parameters

figure()
plt.grid()

cut = goodZspecList
plot(np.arange(1,137), pfitGainShiftVsVoltageList[:,0], '.k')
plot(np.arange(1,137)[cut], pfitGainShiftVsVoltageList[cut,0], '.r')
plt.xlabel('Detector Number')
plt.ylabel("PMT Coefficient")
plt.ylim([0,0.04])


# Examine expected voltage increment for given gain shift
gs = 1.05
vIncrement1 = gainCorrectionCoefficient * log(gs)
vIncrement2 = (np.power(gs, 1/gainCorrectionCoefficient2) - 1.0) * voltageSet

plt.figure()
plot(np.arange(1, 137), vIncrement2, '.k')
cutt = goodZspecList
plot(np.arange(1, 137)[cutt], vIncrement2[cutt], '.r')

# SCATTER Plot of the two v increment estimates
plt.figure()
plt.grid()
plot(vIncrement1[goodZspecList], vIncrement2[goodZspecList], '.k')

plt.xlim([0, 3.5])
plt.ylim([0, 3.5])

plt.xlabel('V Increment 1')
plt.ylabel('V Increment 2')

# SCATTER Plot of the ratio of the increment estimates
plt.figure()
plt.grid()
plot(vIncrement1[goodZspecList], vIncrement2[goodZspecList]/vIncrement1[goodZspecList], '.k')

plt.xlim([0, 3.5])
plt.ylim([0, 3.5])

plt.xlabel('V Increment 1')
plt.ylabel('V Increment 2 / V Increment 1')



# Different in v-increment vs detector number
plt.figure()
plt.grid()
plot(np.arange(1, 137)[goodZspecList], vIncrement2[goodZspecList] - vIncrement1[goodZspecList], '.k')

#plt.xlim([0, 3.5])
#plt.ylim([0, 3.5])

plt.xlabel('Detector Number')
plt.ylabel('V increment 2 - V increment 1')